package com.baeldung.libraries.debezium.config;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.runtime.WorkerConfig;
import org.apache.kafka.connect.runtime.standalone.StandaloneConfig;
import org.apache.kafka.connect.storage.MemoryOffsetBackingStore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.NoSuchFileException;
import java.util.HashMap;
import java.util.Map;

public class JacksonFileOffsetBackingStore extends MemoryOffsetBackingStore {

    private static final ObjectMapper objectMapper = new ObjectMapper();

    private static final Logger log = LoggerFactory.getLogger(JacksonFileOffsetBackingStore.class);

    private File file;

    public JacksonFileOffsetBackingStore() {
    }

    @Override
    public void configure(WorkerConfig config) {
        super.configure(config);
        file = new File(config.getString(StandaloneConfig.OFFSET_STORAGE_FILE_FILENAME_CONFIG));
    }

    @Override
    public synchronized void start() {
        super.start();
        log.info("Starting JacksonFileOffsetBackingStore with file {}", file);
        load();
    }

    @Override
    public synchronized void stop() {
        super.stop();
        // Nothing to do since this doesn't maintain any outstanding connections/data
        log.info("Stopped JacksonFileOffsetBackingStore");
    }

    @SuppressWarnings("unchecked")
    private void load() {
        try {
            Map<byte[], byte[]> raw = objectMapper.readValue(file, new TypeReference<>() {});
            data = new HashMap<>();
            for (Map.Entry<byte[], byte[]> mapEntry : raw.entrySet()) {
                ByteBuffer key = (mapEntry.getKey() != null) ? ByteBuffer.wrap(mapEntry.getKey()) : null;
                ByteBuffer value = (mapEntry.getValue() != null) ? ByteBuffer.wrap(mapEntry.getValue()) : null;
                data.put(key, value);
            }
        } catch (NoSuchFileException | FileNotFoundException e) {
            // NoSuchFileException: Ignore, may be new.
        } catch (IOException e) {
            throw new ConnectException(e);
        }
    }

    @Override
    protected void save() {
        try {
            Map<byte[], byte[]> raw = new HashMap<>();
            for (Map.Entry<ByteBuffer, ByteBuffer> mapEntry : data.entrySet()) {
                byte[] key = (mapEntry.getKey() != null) ? mapEntry.getKey().array() : null;
                byte[] value = (mapEntry.getValue() != null) ? mapEntry.getValue().array() : null;
                raw.put(key, value);
            }
            objectMapper.writeValue(file, raw);
        } catch (IOException e) {
            throw new ConnectException(e);
        }
    }
}
